import matplotlib.pyplot as plt
from random_walk import RandomWalk
import math

plt.style.use('classic')
plt.rcParams['font.sans-serif'] = ['SimHei']  # 显示中文标签
plt.rcParams['axes.unicode_minus'] = False
while True:
    # 创建一个RandomWalk实例。
    rw = RandomWalk(50000)
    rw.fill_walk()
    # 将所有的点都绘制出来。
    fig, ax = plt.subplots(figsize=(10,6), dpi=128)
    ax.scatter(rw.x_values, rw.y_values, s=1,
               c=range(rw.num_points),
               # c=[rw.x_values[x] ** 2 + rw.y_values[x] ** 2 for x in range(rw.num_points)],
               cmap=plt.cm.jet,
               linewidths=0)
    # 突出起点和终点。
    ax.scatter(0, 0, c='orangered', edgecolors='red', s=400,marker='^')
    ax.scatter(rw.x_values[-1], rw.y_values[-1], c='lime', edgecolors='lime',
               s=400,marker='v')

    # 隐藏坐标轴。
    ax.get_xaxis().set_visible(False)
    ax.get_yaxis().set_visible(False)

    plt.show()

    keep_running = input("Make another walk? (y/n): ")
    if keep_running == 'n':
        break
